package com.idatax.auth.common.jwt;

import com.nimbusds.jose.JOSEException;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.jwk.Curve;
import com.nimbusds.jose.jwk.ECKey;
import com.nimbusds.jose.jwk.JWK;
import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.KeyUse;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jose.jwk.gen.ECKeyGenerator;
import com.nimbusds.jose.jwk.gen.RSAKeyGenerator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import javax.annotation.PostConstruct;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Service
public class JwkKeyService {
    private static final Logger log = LoggerFactory.getLogger(JwkKeyService.class);
    private final Map<String, JWK> keys = new ConcurrentHashMap<>();
    private JWK currentSigningKey;
    private String currentSigningKeyId;

    @Value("${jwt.key.algorithm:RS256}")
    private String keyAlgorithm;

    @PostConstruct
    public void init() {
        // Clean up algorithm config value
        if(keyAlgorithm != null) {
            keyAlgorithm = keyAlgorithm.split("#")[0].trim();
        }
        // Initialize keys after injection
        if (keyAlgorithm != null) {
            if (keyAlgorithm.startsWith("RS")) {
                generateRsaKey();
            } else {
                generateEcKey();
            }
        }
    }

    @Value("${jwt.key.rotation.hours:24}")
    private int rotationHours;

    public JwkKeyService() {
        // Initialization moved to init() after @Value injection
    }

    public JWKSet getJwkSet() {
        return new JWKSet(new ArrayList<>(keys.values()));
    }

    public JWK getSigningKey() {
        if (currentSigningKey == null) {
            String errorMsg = String.format("No valid signing key found for algorithm %s", keyAlgorithm);
            log.error(errorMsg);
            throw new RuntimeException(errorMsg);
        }
        log.debug("Using signing key with ID: {}", currentSigningKeyId);
        return currentSigningKey;
    }

    @Scheduled(fixedRateString = "${jwt.key.rotation.hours:24}", timeUnit = TimeUnit.HOURS)
    public void scheduledKeyRotation() {
        if ("RSA".equals(keyAlgorithm)) {
            generateRsaKey();
        } else {
            generateEcKey();
        }
    }

    public void generateRsaKey() {
        try {
            RSAKey rsaKey = new RSAKeyGenerator(2048)
                    .keyUse(KeyUse.SIGNATURE)
                    .algorithm(JWSAlgorithm.RS256)
                    .keyID(UUID.randomUUID().toString())
                    .generate();
            keys.put(rsaKey.getKeyID(), rsaKey);
            currentSigningKey = rsaKey;
            currentSigningKeyId = rsaKey.getKeyID();
            log.info("Generated new RSA key with ID: {}", rsaKey.getKeyID());
        } catch (JOSEException e) {
            log.error("Failed to generate RSA key", e);
            throw new RuntimeException("Failed to generate RSA key", e);
        }
    }

    public void generateEcKey() {
        try {
            JWSAlgorithm algorithm = new JWSAlgorithm(keyAlgorithm);
            Curve curve = getCurveForAlgorithm(algorithm);
            
            ECKey ecKey = new ECKeyGenerator(curve)
                    .keyUse(KeyUse.SIGNATURE)
                    .algorithm(algorithm)
                    .keyID(UUID.randomUUID().toString())
                    .generate();
            keys.put(ecKey.getKeyID(), ecKey);
            currentSigningKey = ecKey;
            currentSigningKeyId = ecKey.getKeyID();
            log.info("Generated new EC key with ID: {}", ecKey.getKeyID());
        } catch (JOSEException e) {
            log.error("Failed to generate EC key", e);
            throw new RuntimeException("Failed to generate EC key", e);
        }
    }

    public void rotateKeys() {
        if (keyAlgorithm.startsWith("RS")) {
            generateRsaKey();
        } else {
            generateEcKey();
        }
    }

    private Curve getCurveForAlgorithm(JWSAlgorithm algorithm) {
        switch (algorithm.getName()) {
            case "ES256": return Curve.P_256;
            case "ES384": return Curve.P_384;
            case "ES512": return Curve.P_521;
            default: throw new IllegalArgumentException("Unsupported EC algorithm: " + algorithm);
        }
    }
}
